import json
import os
from glob import glob
import numpy as np
import cv2
from sklearn.utils import shuffle

def read_labelme(json_file):
    f = open(json_file, "r")
    data = json.load(f)
    f.close()
    shapes = data["shapes"]
    img_name=data['imagePath']
    bboxes = []
    labels=[]
    for shape in shapes:
        pt = np.array(shape['points'], dtype=np.int)
        if len(pt) == 2:
            pt =np.array([pt[0], [pt[1][0], pt[0][1]], pt[1], [pt[0][0], pt[1][1]]])
        bboxes.append(pt)
        labels.append(shape['label'])
    assert len(labels)==len(bboxes)
    info=dict(
        bbox=bboxes,
        label=labels,
        path=img_name
    )
    return info

if __name__=='__main__':
    data_dirs=['/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/DataSet/芜湖海螺/水泥船/height/20220105',
               '/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/DataSet/芜湖海螺/水泥船/height/20220107']
    dst_dir='/media/wsl/a9f0161f-7971-c843-8c81-c68049a0235a/DataSet/芜湖海螺/水泥船/高度/20220106'
    dst_img_train_dir=os.path.join(dst_dir,'img_dir','train')
    dst_img_val_dir = os.path.join(dst_dir, 'img_dir', 'val')
    dst_anno_train_dir = os.path.join(dst_dir, 'anno_dir', 'train')
    dst_anno_val_dir = os.path.join(dst_dir, 'anno_dir', 'val')
    os.makedirs(dst_anno_val_dir,exist_ok=True)
    os.makedirs(dst_anno_train_dir,exist_ok=True)
    os.makedirs(dst_img_val_dir,exist_ok=True)
    os.makedirs(dst_img_train_dir,exist_ok=True)

    train_rate=0.8
    cls=['0']
    one_flag=True
    idx=0
    for img_dir in data_dirs:
        json_files = glob(img_dir + '/*.json')
        json_files=shuffle(json_files)
        num=len(json_files)
        train_files=json_files[0:int(num*train_rate)]
        val_files=json_files[int(num*train_rate):]

        for json_file in train_files:
            info=read_labelme(json_file)
            bbox=info['bbox']
            label=info['label']
            img_name=info['path']
            if len(bbox) == 0:
                continue
            im_path = os.path.join(img_dir, img_name)
            im=cv2.imread(im_path)
            cv2.imwrite(os.path.join(dst_img_train_dir,'%d.png'%idx),im)
            mask=np.zeros(im.shape[0:2],dtype=np.uint8)
            if one_flag:
                mask = cv2.fillPoly(mask,bbox,color=1)
            else:
                for c,bb in zip(label,bbox):
                    mask = cv2.fillPoly(mask, bbox, color=cls.index(c)+1)
            cv2.imwrite(os.path.join(dst_anno_train_dir,'%d.png'%idx),mask)
            idx+=1

        for json_file in val_files:
            info=read_labelme(json_file)
            bbox=info['bbox']
            label=info['label']
            img_name=info['path']
            if len(bbox) == 0:
                continue
            im_path = os.path.join(img_dir, img_name)
            im=cv2.imread(im_path)
            cv2.imwrite(os.path.join(dst_img_val_dir,'%d.png'%idx),im)
            mask=np.zeros(im.shape[0:2],dtype=np.uint8)
            if one_flag:
                mask = cv2.fillPoly(mask,bbox,color=1)
            else:
                for c,bb in zip(label,bbox):
                    mask = cv2.fillPoly(mask, bbox, color=cls.index(c)+1)
            cv2.imwrite(os.path.join(dst_anno_val_dir,'%d.png'%idx),mask)
            idx+=1

